from tpot import TPOTClassifier
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction import DictVectorizer

# 数据加载
df = pd.read_csv('./train.csv')

# 数据探索
# 查看train_data信息
#pd.set_option('display.max_columns', None) #显示所有列
print('查看数据信息：列名、非空个数、类型等')
print(df.info())
print('-'*30)
print('查看数据摘要')
print(df.describe())
print('-'*30)
print('查看离散数据分布')
print(df.describe(include=['O']))
print('-'*30)
print('查看前5条数据')
print(df.head())
print('-'*30)
print('查看后5条数据')
print(df.tail())

# 使用平均年龄来填充年龄中的nan值
df['Age'].fillna(df['Age'].mean(), inplace=True)

#print(df['Embarked'].value_counts())
# 使用登录最多的港口来填充登录港口的nan值
df['Embarked'].fillna('S', inplace=True)

# # 特征选择
features = ['Pclass', 'Sex', 'Age', 'SibSp', 'Parch', 'Fare', 'Embarked']
X = df[features]
y = df['Survived']

dvec = DictVectorizer(sparse=False)
X = dvec.fit_transform(X.to_dict(orient='record'))

X_train, X_test, y_train, y_test = train_test_split(X, y,
                                                    train_size=0.75, test_size=0.25)

pipeline_optimizer = TPOTClassifier(generations=5, population_size=20, cv=5,
                                    random_state=42, verbosity=2)
pipeline_optimizer.fit(X_train, y_train)
print(pipeline_optimizer.score(X_test, y_test))
pipeline_optimizer.export('tpot_exported_titanic.py')